#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
API文档生成脚本

该脚本用于自动生成和更新SQL Agent API的RESTful风格接口文档。
基于FastAPI应用的路由和模型定义，自动提取API信息并生成Markdown格式的文档。

使用方法:
    python scripts/generate_api_docs.py

作者: SQL Agent Team
版本: 1.0.0
"""

import os
import sys
import json
import inspect
from datetime import datetime
from typing import Dict, List, Any
from pathlib import Path

# 添加项目根目录到Python路径
project_root = Path(__file__).parent.parent
sys.path.insert(0, str(project_root))

try:
    from fastapi import FastAPI
    from fastapi.routing import APIRoute
    from pydantic import BaseModel
    
    # 尝试导入API模块，如果失败则使用模拟数据
    try:
        from src.api.models import (
            SqlQueryRequest, SqlQueryResponse,
            ChatCompletionRequest, ChatCompletionResponse,
            Message
        )
        MODELS_AVAILABLE = True
    except ImportError:
        print("警告: 无法导入API模型，将使用模拟数据生成文档")
        MODELS_AVAILABLE = False
        
        # 创建模拟模型类
        class Message(BaseModel):
            role: str
            content: str
            
        class SqlQueryRequest(BaseModel):
            query: str
            userId: str = None
            conversationId: str = None
            database_uri: str = "sqlite:///Chinook.db"
            
        class SqlQueryResponse(BaseModel):
            id: str
            result: str
            success: bool = True
            error: str = None
            created: int
            
        class ChatCompletionRequest(BaseModel):
            messages: List[dict]
            stream: bool = False
            userId: str = None
            conversationId: str = None
            
        class ChatCompletionResponse(BaseModel):
            id: str
            object: str = "chat.completion"
            created: int
            choices: List[dict]
            
except ImportError as e:
    print(f"导入错误: {e}")
    print("请确保已安装FastAPI和Pydantic")
    sys.exit(1)


class APIDocGenerator:
    """API文档生成器"""
    
    def __init__(self, fastapi_app: FastAPI = None):
        self.app = fastapi_app
        self.base_url = "http://localhost:8013"
        self.version = "1.0.0"
        
        # 静态路由信息（当无法导入FastAPI应用时使用）
        self.static_routes = [
            {
                'path': '/',
                'methods': ['GET'],
                'name': '根端点',
                'description': '获取API基本信息',
                'docstring': '根路径，返回API信息。'
            },
            {
                'path': '/health',
                'methods': ['GET'],
                'name': '健康检查',
                'description': '检查服务健康状态',
                'docstring': '健康检查端点。'
            },
            {
                'path': '/sql/query',
                'methods': ['POST'],
                'name': 'SQL查询',
                'description': '执行SQL查询或回答数据库相关问题',
                'docstring': '执行SQL查询或回答数据库相关问题。'
            },
            {
                'path': '/sql/health',
                'methods': ['GET'],
                'name': 'SQL健康检查',
                'description': 'SQL模块的健康检查',
                'docstring': 'SQL模块的健康检查端点。'
            },
            {
                'path': '/v1/chat/completions',
                'methods': ['POST'],
                'name': '聊天完成接口',
                'description': '提供与OpenAI ChatGPT API兼容的聊天完成接口',
                'docstring': '兼容OpenAI格式的聊天完成接口，支持流式和非流式响应。'
            }
        ]
        
    def extract_routes(self) -> List[Dict[str, Any]]:
        """提取FastAPI应用的所有路由信息"""
        if self.app is None:
            # 如果没有FastAPI应用，使用静态路由信息
            return self.static_routes
            
        routes = []
        
        try:
            for route in self.app.routes:
                if isinstance(route, APIRoute):
                    route_info = {
                        'path': route.path,
                        'methods': list(route.methods),
                        'name': route.name,
                        'summary': getattr(route, 'summary', ''),
                        'description': getattr(route, 'description', ''),
                        'tags': getattr(route, 'tags', []),
                        'endpoint': route.endpoint,
                    }
                    
                    # 提取函数文档字符串
                    if route.endpoint and hasattr(route.endpoint, '__doc__'):
                        route_info['docstring'] = route.endpoint.__doc__ or ''
                    
                    # 提取参数信息
                    if hasattr(route.endpoint, '__annotations__'):
                        route_info['annotations'] = route.endpoint.__annotations__
                    
                    routes.append(route_info)
        except Exception as e:
            print(f"警告: 提取路由信息时出错，使用静态数据: {e}")
            return self.static_routes
                
        return routes if routes else self.static_routes
    
    def get_model_schema(self, model_class: BaseModel) -> Dict[str, Any]:
        """获取Pydantic模型的JSON Schema"""
        try:
            return model_class.model_json_schema()
        except AttributeError:
            # 兼容旧版本Pydantic
            return model_class.schema()
    
    def format_model_example(self, model_class: BaseModel) -> str:
        """生成模型的JSON示例"""
        try:
            schema = self.get_model_schema(model_class)
            example = {}
            
            if 'properties' in schema:
                for field_name, field_info in schema['properties'].items():
                    field_type = field_info.get('type', 'string')
                    
                    if field_type == 'string':
                        if 'example' in field_info:
                            example[field_name] = field_info['example']
                        elif field_name.lower() in ['id', 'userid', 'conversationid']:
                            example[field_name] = f"example_{field_name.lower()}"
                        elif field_name.lower() == 'query':
                            example[field_name] = "查询示例"
                        else:
                            example[field_name] = f"示例{field_name}"
                    elif field_type == 'integer':
                        example[field_name] = 1234567890 if 'time' in field_name.lower() else 0
                    elif field_type == 'boolean':
                        example[field_name] = True
                    elif field_type == 'array':
                        example[field_name] = []
                    else:
                        example[field_name] = None
            
            return json.dumps(example, ensure_ascii=False, indent=2)
        except Exception as e:
            return f"{{\n  \"error\": \"无法生成示例: {str(e)}\"\n}}"
    
    def generate_endpoint_doc(self, route: Dict[str, Any]) -> str:
        """生成单个端点的文档"""
        doc = []
        
        # 端点标题
        methods_str = "/".join(sorted(route['methods']))
        doc.append(f"#### {methods_str} {route['path']}")
        doc.append("")
        
        # 描述
        description = route.get('docstring', '').strip()
        if description:
            # 提取第一行作为简短描述
            first_line = description.split('\n')[0].strip()
            doc.append(first_line)
            doc.append("")
        
        # 请求示例
        if 'POST' in route['methods']:
            doc.append("**请求示例:**")
            doc.append("```bash")
            doc.append(f'curl -X POST "{self.base_url}{route["path"]}" \\')
            doc.append('  -H "Content-Type: application/json" \\')
            doc.append('  -d \'{')
            doc.append('    "key": "value"')
            doc.append('  }\'')
            doc.append("```")
            doc.append("")
        elif 'GET' in route['methods']:
            doc.append("**请求示例:**")
            doc.append("```bash")
            doc.append(f'curl -X GET "{self.base_url}{route["path"]}"')
            doc.append("```")
            doc.append("")
        
        # 响应示例
        doc.append("**响应示例:**")
        doc.append("```json")
        doc.append("{")
        doc.append('  "status": "success",')
        doc.append('  "data": {}')
        doc.append("}")
        doc.append("```")
        doc.append("")
        doc.append("---")
        doc.append("")
        
        return "\n".join(doc)
    
    def generate_models_doc(self) -> str:
        """生成数据模型文档"""
        doc = []
        doc.append("## 数据模型")
        doc.append("")
        
        models = [
            ("Message", Message),
            ("SqlQueryRequest", SqlQueryRequest),
            ("SqlQueryResponse", SqlQueryResponse),
            ("ChatCompletionRequest", ChatCompletionRequest),
            ("ChatCompletionResponse", ChatCompletionResponse),
        ]
        
        for model_name, model_class in models:
            doc.append(f"### {model_name}")
            doc.append("```json")
            doc.append(self.format_model_example(model_class))
            doc.append("```")
            doc.append("")
        
        return "\n".join(doc)
    
    def generate_full_documentation(self) -> str:
        """生成完整的API文档"""
        doc = []
        
        # 文档头部
        doc.append("# SQL Agent API 接口文档")
        doc.append("")
        doc.append(f"**生成时间**: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")
        doc.append(f"**API版本**: {self.version}")
        doc.append(f"**基础URL**: {self.base_url}")
        doc.append("")
        
        # API概述
        doc.append("## 概述")
        doc.append("")
        doc.append("SQL Agent API 是一个基于 LangGraph 的智能 SQL 查询服务，")
        doc.append("提供自然语言到 SQL 的转换和数据库查询功能。")
        doc.append("")
        
        # 提取并生成端点文档
        routes = self.extract_routes()
        
        doc.append("## API 端点")
        doc.append("")
        
        for i, route in enumerate(routes, 1):
            doc.append(f"### {i}. {route.get('name', '未命名端点')}")
            doc.append("")
            doc.append(self.generate_endpoint_doc(route))
        
        # 添加数据模型文档
        doc.append(self.generate_models_doc())
        
        # 使用示例
        doc.append("## 使用示例")
        doc.append("")
        doc.append("### Python 示例")
        doc.append("")
        doc.append("```python")
        doc.append("import requests")
        doc.append("")
        doc.append(f'base_url = "{self.base_url}"')
        doc.append("")
        doc.append("# 健康检查")
        doc.append('response = requests.get(f"{base_url}/health")')
        doc.append('print("健康检查:", response.json())')
        doc.append("")
        doc.append("# SQL查询")
        doc.append("sql_query = {")
        doc.append('    "query": "查询前5个艺术家的信息",')
        doc.append('    "userId": "user123",')
        doc.append('    "conversationId": "conv456"')
        doc.append("}")
        doc.append('response = requests.post(f"{base_url}/sql/query", json=sql_query)')
        doc.append('print("SQL查询结果:", response.json())')
        doc.append("```")
        doc.append("")
        
        # 错误码
        doc.append("## 错误码")
        doc.append("")
        doc.append("| HTTP状态码 | 错误类型 | 描述 |")
        doc.append("|------------|----------|------|")
        doc.append("| 200 | 成功 | 请求成功处理 |")
        doc.append("| 400 | 请求错误 | 请求参数有误 |")
        doc.append("| 500 | 服务器错误 | 服务器内部错误 |")
        doc.append("| 503 | 服务不可用 | 服务未初始化或不可用 |")
        doc.append("")
        
        # 注意事项
        doc.append("## 注意事项")
        doc.append("")
        doc.append("1. **数据库连接**: 默认使用 SQLite 数据库 `Chinook.db`")
        doc.append("2. **会话管理**: 使用 `userId` 和 `conversationId` 来维护会话上下文")
        doc.append("3. **流式响应**: 流式响应使用 Server-Sent Events (SSE) 格式")
        doc.append("4. **错误处理**: 建议客户端检查 `success` 字段和 HTTP 状态码")
        doc.append("")
        
        return "\n".join(doc)
    
    def save_documentation(self, output_path: str = None) -> str:
        """保存文档到文件"""
        if output_path is None:
            output_path = project_root / "docs" / "API_Documentation_Generated.md"
        
        # 确保输出目录存在
        output_path = Path(output_path)
        output_path.parent.mkdir(parents=True, exist_ok=True)
        
        # 生成文档内容
        content = self.generate_full_documentation()
        
        # 写入文件
        with open(output_path, 'w', encoding='utf-8') as f:
            f.write(content)
        
        return str(output_path)


def main():
    """主函数"""
    print("开始生成API文档...")
    
    try:
        # 尝试导入FastAPI应用
        app = None
        try:
            from src.api.main import app as fastapi_app
            app = fastapi_app
            print("成功导入FastAPI应用")
        except ImportError:
            print("无法导入FastAPI应用，将使用静态数据生成文档")
        
        # 创建文档生成器
        generator = APIDocGenerator(app)
        
        # 生成并保存文档
        output_path = generator.save_documentation()
        
        print(f"API文档已生成: {output_path}")
        print("\n文档包含以下内容:")
        print("- API概述和基本信息")
        print("- 所有端点的详细说明")
        print("- 请求/响应示例")
        print("- 数据模型定义")
        print("- 使用示例")
        print("- 错误码说明")
        
    except Exception as e:
        print(f"生成文档时出错: {str(e)}")
        import traceback
        traceback.print_exc()
        sys.exit(1)


if __name__ == "__main__":
    main()